import torch
from torch import nn
import torchvision
import os
import struct
import onnxruntime as ort
import numpy as np
import struct

def main():
    print('cuda device count: ', torch.cuda.device_count())
    ort_session = ort.InferenceSession('yolov5s.onnx')
    #mydata=np.random.random((1,3,384,640)).astype(np.float32)
    mydata=np.ones((1,3,640,640)).astype(np.float32)

    #mydata.tofile("mydata.bin")
    outputs = ort_session.run(None, {'images': mydata})
    result0 = outputs[0][0].transpose(0,1,3,2).reshape(-1);
    result1 = outputs[1][0].transpose(0,1,3,2).reshape(-1);
    result2 = outputs[2][0].transpose(0,1,3,2).reshape(-1);
    print(result0[0:10])
    print(result1[0:10])
    print(result2[0:10])
    #print(outputs[2][0][0][0][0:5])

if __name__ == '__main__':
    main()

